/*
 * @Author       : mark
 * @Date         : 2020-06-15
 * @copyleft Apache 2.0
 */

#ifndef THREADPOOL_H
#define THREADPOOL_H

#include <vector>
#include <queue>
#include <memory>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <functional>
#include <stdexcept>

class ThreadPool
{

public:
    explicit ThreadPool(size_t thread_count = 4)
        : closed(false)
    {
        for (size_t i = 0; i < thread_count; ++i)
            workers.emplace_back(
                std::thread(
                    [this]
                    {
                        for (;;)
                        {
                            std::function<void()> task;
            
                            {
                                std::unique_lock<std::mutex> lock(this->mutex);
                                this->condition.wait(lock,
                                                     [this]
                                                     { return this->closed || this->tasks.size(); });
                                if (this->closed && this->tasks.empty())
                                    return;
                                task = std::move(this->tasks.front());
                                this->tasks.pop();
                            }

                            task();
                        }
                    }));
    }

    ThreadPool() = default;

    ~ThreadPool()
    {
        closed = true;
        condition.notify_all();
        for (auto &t : workers)
            if (t.joinable())
                t.join();
    }

    template <class F, class... Args>
    void AddTask(F&& f, Args &&...args)
    {
        {
            std::unique_lock<std::mutex> lock(mutex);
            tasks.emplace(f);
        }
        condition.notify_one();
    }
private:
    using TASK = std::function<void()>;
    std::mutex mutex;
    std::queue<TASK> tasks;
    std::condition_variable condition;
    std::vector<std::thread> workers;
    std::atomic<bool> closed;
};

#endif // THREADPOOL_H